import dask.dataframe as dd
import pandas as pd
import sqlite3
from dask.distributed import Client
from concurrent.futures import ThreadPoolExecutor
import os
from time import time
import warnings

# 忽略关于默认样式的警告
warnings.filterwarnings('ignore', category=UserWarning, message='Workbook contains no default style, apply openpyxl\'s default')
warnings.filterwarnings('ignore', category=UserWarning, message='Port 8787 is already in use.Perhaps you already have a cluster running?')


def clean_data(df: pd.DataFrame) -> pd.DataFrame:
    """
    清理DataFrame中的空值，删除任何包含空值的行。
    """
    # 删除任何包含空值的行
    print("清理DataFrame中的空值，删除任何包含空值的行。")
    df_cleaned = df.dropna()

    return df_cleaned


def read_and_clean_file(file_path: str) -> pd.DataFrame:
    """
    读取Excel文件并清理数据。
    """
    print("读取Excel文件并清理数据。")
    df = pd.read_excel(file_path)
    df_cleaned = clean_data(df)
    return df_cleaned


def main():
    file_path = r'D:\文件\大四\毕业论文\数据'
    file_names = ['TRD_Dalyr', 'TRD_Dalyr1', 'TRD_Dalyr2', 'TRD_Dalyr3', 'TRD_Dalyr4', 'TRD_Dalyr5']

    print("创建文件路径映射字典")
    file_mapper = {name: os.path.join(file_path, f"{name}.xlsx") for name in file_names}

    print("初始化Dask客户端")
    client = Client()

    print("使用多线程读取和清理文件")
    with ThreadPoolExecutor() as executor:
        print("提交任务到线程池")
        futures = [executor.submit(read_and_clean_file, full_path) for full_path in file_mapper.values()]

        print("获取清理后的DataFrame列表")
        cleaned_dfs = []
        for index, future in enumerate(futures, start=1):
            cleaned_dfs.append(future.result())
            print(f"已处理文件 {index}/{len(futures)}")

    print("# 将清理后的DataFrame列表合并为Dask DataFrame")
    ddf = dd.concat([dd.from_pandas(df, npartitions=1) for df in cleaned_dfs], axis=0)

    print("# 将Dask DataFrame保存为SQLite数据库")
    database_path = os.path.join(file_path, 'combined_data.db')
    ddf.to_sql('combined_data', f'sqlite:///{database_path}', if_exists='replace', index=False)

    # 从SQLite数据库读取数据到Dask DataFrame（可选）
    # ddf_from_db = dd.read_sql_table('combined_data', sqlite3.connect(database_path), index_col='index')

    # 打印结果（可选）
    # print(ddf_from_db.head())

    # 关闭Dask客户端
    client.close()


if __name__ == "__main__":
    main()
